import os
import pandas as pd
import pymysql

# 数据库连接配置
db_config = {
    "host": "???",
    "user": "root",
    "password": "???",
    "db": "industry",
    "charset": "utf8mb4",
    "cursorclass": pymysql.cursors.DictCursor,
}


def get_new_columns(df, order: int):
    columns = df.columns.to_list()
    start = columns.index("网址") + 1
    end = columns.index("其他") + 1
    industry_cls = columns[start:end]
    add_order = lambda x: x + str(order)

    return list(map(add_order, industry_cls))


def set_industry_cls(df, new_columns):
    raw_columns = list(map(lambda x: x[:-1], new_columns))

    def func(item):
        if not isinstance(item, str):
            return {}
        row_data = df[df["企业名称"] == item]
        if len(row_data) > 0:
            row_dict = row_data.iloc[0].to_dict()
        else:
            row_dict = {}

        res = {
            new_k: row_dict.get(raw_k, 0)
            for raw_k, new_k in zip(raw_columns, new_columns)
        }
        return pd.Series(res)

    return func


def get_patent_coporation(input_csv_path):
    output_csv_path = input_csv_path.rsplit(".", 1)[0] + "_output.xlsx"
    if input_csv_path.endswith("xlsx") or input_csv_path.endswith("xls"):
        input_companies = pd.read_excel(input_csv_path)
    elif input_csv_path.endswith("csv"):
        input_companies = pd.read_csv(input_csv_path)

    # 连接到数据库
    connection = pymysql.connect(**db_config)
    try:
        with connection.cursor() as cursor:
            # 准备一个空的列表来存储所有结果
            all_results = []
            # 准备一个集合来跟踪唯一的（企业名称, 合作企业名称）对
            company_pairs = set()
            # 定义每批处理的企业名称数量
            batch_size = 1000
            # 分批次遍历DataFrame
            for start in range(0, len(input_companies), batch_size):
                end = start + batch_size
                batch_group = input_companies.iloc[start:end]
                # 确保企业名称是字符串类型，并转换为元组
                company_names = tuple(batch_group["企业名称"].astype(str))
                # 构建SQL查询
                placeholders = ", ".join(["%s"] * len(company_names))
                query_sql = f"""
                SELECT pc.company_one AS 企业名称,
                    pc.company_two AS 合作企业名称,
                    pc.company_one_longitude AS 企业经度,
                    pc.company_one_latitude AS 企业纬度,
                    pc.company_two_longitude AS 合作企业经度,
                    pc.company_two_latitude AS 合作企业纬度,
                    pc.cooperate_number AS 合作次数
                FROM patent_cooperate pc
                WHERE pc.company_one IN ({placeholders}) OR pc.company_two IN ({placeholders})
                """
                # 执行查询
                cursor.execute(query_sql, company_names * 2)
                # 获取查询结果
                result_rows = cursor.fetchall()
                # 遍历查询结果，并检查是否已存在唯一的（企业名称, 合作企业名称）对
                for row in result_rows:
                    # 将企业名称和合作企业名称排序，以确保顺序一致性
                    sorted_pair = tuple(sorted([row["企业名称"], row["合作企业名称"]]))
                    # 如果这个排序后的对不在company_pairs集合中，则添加到all_results和company_pairs
                    if sorted_pair not in company_pairs:
                        all_results.append(row)
                        company_pairs.add(sorted_pair)

        # 将查询结果转换为DataFrame
        df = pd.DataFrame(
            all_results,
            columns=[
                "企业名称",
                "合作企业名称",
                "企业经度",
                "企业纬度",
                "合作企业经度",
                "合作企业纬度",
                "合作次数",
            ],
        )
        if len(df) == 0:
            return
        
        add_columns1 = get_new_columns(input_companies, 1)
        add_columns2 = get_new_columns(input_companies, 2)
        df[add_columns1] = df["企业名称"].apply(
            set_industry_cls(input_companies, add_columns1)
        )
        df[add_columns2] = df["合作企业名称"].apply(
            set_industry_cls(input_companies, add_columns2)
        )

        df.to_excel(output_csv_path, index=False, header=True)

    finally:
        # 关闭数据库连接
        connection.close()
    print(f"Query results have been written to {output_csv_path}")


if __name__ == "__main__":
    folder = "/home/jie/gitee/pku_industry/csv_folder_test"
    for file in os.listdir(folder):
        if file.endswith("专利统计.csv"):
            print(file)
            file_name = os.path.join(folder, file)
            get_patent_coporation(file_name)
            print()